#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试数据库初始化脚本
创建包含默认数据的数据库，用于开发和测试
"""

import sys
import os
from pathlib import Path

# 添加项目根目录到Python路径
current_dir = Path(__file__).parent.parent
sys.path.insert(0, str(current_dir))

from src.core.path_manager import path_manager
from src.models.database.database import DatabaseManager

def init_test_database():
    """初始化测试数据库"""
    print("========================================")
    print("WMS Sample 测试数据库初始化")
    print("========================================")
    print()
    
    # 获取数据库路径
    db_path = path_manager.get_database_file()
    print(f"数据库路径: {db_path}")
    print(f"部署模式: {path_manager.deployment_mode}")
    print()
    
    # 检查是否已存在数据库
    if db_path.exists():
        print("⚠️  数据库文件已存在")
        choice = input("是否要删除现有数据库并创建新的测试数据库？(y/N): ").strip().lower()
        if choice != 'y':
            print("操作已取消")
            return
        
        # 删除现有数据库
        try:
            os.remove(str(db_path))
            print("✅ 已删除现有数据库文件")
        except Exception as e:
            print(f"❌ 删除数据库文件失败: {e}")
            return
    
    print("正在创建测试数据库...")
    
    try:
        # 创建包含默认数据的数据库管理器
        test_db_manager = DatabaseManager(init_data=True)
        
        print("✅ 数据库表结构创建完成")
        print("✅ 默认数据初始化完成")
        
        # 验证数据
        session = test_db_manager.get_session()
        try:
            from src.models.database.models import Plate, Position
            
            # 检查数据
            plates = session.query(Plate).all()
            total_positions = session.query(Position).count()
            
            print(f"✅ 创建了 {len(plates)} 个料盘:")
            for plate in plates:
                positions_count = session.query(Position).filter_by(plate_id=plate.id).count()
                print(f"   - {plate.name}: 容量 {plate.capacity}, 料位 {positions_count}")
            
            print(f"✅ 总计 {total_positions} 个料位")
            
            print("🎉 测试数据库创建成功！")
            print("   - 包含完整的表结构")
            print("   - 包含任务追踪字段")
            print("   - 包含默认料盘和料位")
            print("   - 适合开发和测试")
                
        finally:
            session.close()
            
    except Exception as e:
        print(f"❌ 创建测试数据库失败: {e}")
        import traceback
        traceback.print_exc()
        return
    
    print()
    print("========================================")
    print("数据库初始化完成")
    print("========================================")
    print()
    print("使用说明:")
    print("1. 此数据库包含默认的料盘和料位数据")
    print("2. 适合用于开发和测试环境")
    print("3. 如需干净数据库，请运行 init_clean_database.py")

def add_sample_data():
    """添加一些示例数据用于测试"""
    print("正在添加示例数据...")
    
    try:
        from src.models.database.database import db_manager
        from src.models.database.models import Position
        from datetime import datetime
        
        session = db_manager.get_session()
        try:
            # 获取一些料位来放置示例样品
            available_positions = session.query(Position).filter_by(status="available").limit(5).all()
            
            if not available_positions:
                print("没有可用的料位来放置示例数据")
                return
            
            # 添加一些示例样品
            sample_data = [
                {"sample_id": "TEST001", "task_type": "place", "operator": "admin"},
                {"sample_id": "TEST002", "task_type": "place", "operator": "admin"},
                {"sample_id": "TEST003", "task_type": "place", "operator": "admin"},
            ]
            
            current_time = datetime.now().isoformat()
            
            for i, sample in enumerate(sample_data):
                if i < len(available_positions):
                    position = available_positions[i]
                    position.sample_id = sample["sample_id"]
                    position.status = "occupied"
                    position.task_id = f"TASK_{i+1:03d}"
                    position.task_type = sample["task_type"]
                    position.operator = sample["operator"]
                    position.operation_time = current_time
            
            session.commit()
            print(f"✅ 添加了 {len(sample_data)} 个示例样品")
            
        finally:
            session.close()
            
    except Exception as e:
        print(f"❌ 添加示例数据失败: {e}")

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--with-samples":
        init_test_database()
        add_sample_data()
    else:
        init_test_database() 